import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('TkAgg')

# 加载数据
X = np.load('./X.npy')  # 注意：大写 X
y = np.load('./y.npy')

# 划分训练集和测试集
x_train, x_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

acc = []
depths = range(2, 50)  # 不需要太大，树太深容易过拟合

for n in depths:
    tree = DecisionTreeClassifier(max_depth=n, random_state=42)
    tree.fit(x_train, y_train)
    acc.append(tree.score(x_test, y_test))

# 可视化
plt.figure(figsize=(10, 6))
plt.plot(depths, acc, 'b-o')
plt.xlabel('Max Depth')
plt.ylabel('Accuracy')
plt.title('Decision Tree Accuracy vs Max Depth')
plt.grid(True)
plt.xticks(depths[::2])  # 避免刻度太密
plt.tight_layout()
plt.show()

# 输出最佳深度
best_depth = depths[np.argmax(acc)]
print(f"最佳 max_depth: {best_depth}, 测试准确率: {max(acc):.4f}")